#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@文件        :    cifar2image.py
@说明        :    将cifar数据集转化为图像和对应的标签
@时间        :    2021/12/20 17:02:31
@作者        :    钟建平
@版本        :    1.0
'''
import argparse
import pickle
import os
import cv2
from tqdm import tqdm

CIFAR10 = [['data_batch_1', 'data_batch_2', 'data_batch_3', 'data_batch_4', 'data_batch_5'], ['test_batch']]
CIFAR100 = [['train'], ['test']]
CIFAR = {"cifar-10": CIFAR10, "cifar-100": CIFAR100}

def  _parser_args():
    """parse input"""
    parser = argparse.ArgumentParser("convert coco to yolo")
    parser.add_argument('--cifar_path', type=str, required=True, help='the path of cifar')
    parser.add_argument('--save_path', type=str,  required=True,  help='the path ot save data')
    parser.add_argument('--cifar_type', type=str,  default='cifar-10', help='cifar-10 or  cifar-100')
    args, _ = parser.parse_known_args()
    return args


def unpickle(file):
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict

def parse_data(batch_data, save_path, cifar_type):
    file_names = batch_data[b'filenames']
    data = batch_data[b'data']
    if cifar_type == 'cifar-10':
        labels = batch_data[b'labels']
    else:
        labels = batch_data[b'fine_labels']
    for i in tqdm(range(len(file_names))):
        image = data[i]
        image = image.reshape((3, 32, 32))
        image = image.transpose(1, 2, 0)
        file_name = file_names[i]
        label = labels[i]
        image_name = 'c_' + str(label) + '_' + file_name.decode()
        cv2.imwrite(save_path + image_name, image)

def convert(cifar_path, cifar_type, save_path):
    cifar = CIFAR[cifar_type]
    train = cifar[0]
    test = cifar[1]
    if not os.path.exists(save_path):
        os.mkdir(save_path)
    if not os.path.exists(save_path+'train'):
        os.mkdir(save_path+'train')
    if not os.path.exists(save_path+'test'):
        os.mkdir(save_path+'test')
    # parse train
    for batch_file in train:
        batch = unpickle(cifar_path + '/' +batch_file)
        print(batch.keys())
        print(batch[b'fine_labels'])
        print('---------')
        # print(batch[b'batch_label'])
        exit(0)
        parse_data(batch, save_path+'train/', cifar_type)
    # parse test
    for batch_file in test:
        batch = unpickle(cifar_path + '/' +batch_file)
        parse_data(batch, save_path+'test/', cifar_type)

if __name__ == "__main__":
    args  = _parser_args()
    convert(args.cifar_path, args.cifar_type, args.save_path)